import torch

x = torch.tensor([0.2, 0.5, 0.1, 0.9])
print("x =", x)

# 找最大值的索引
idx = torch.argmax(x)
print("最大值索引:", idx.item())
print("最大值:", x[idx].item())

x = torch.tensor([[1.0, 3.0, 2.0],
                  [4.0, 0.5, 5.0]])

print("x =", x)

# 不指定dim，会把x拉平为一维，然后返回最大值在整个tensor中的索引
idx = torch.argmax(x)
print("整体最大值索引:", idx.item())
print("整体最大值:", x.flatten()[idx].item())

x = torch.tensor([[1.0, 3.0, 2.0],
                  [4.0, 0.5, 5.0]])

# dim=1 表示在每一行上求最大值索引
idx = torch.argmax(x, dim=1)
print("每一行最大值索引:", idx)

idx = torch.argmax(x, dim=0)
print("每一列最大值索引:", idx)

# tolist()：把张量（tensor）转换成标准的 Python 列表（list）或嵌套列表结构
print(x[0].tolist())
